
import os
from datetime import datetime

import cv2
import numpy as np

from orbbec.infrared_pipeline import InfraredCamera


if __name__ == "__main__":
    # 创建输出文件夹
    current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
    output_path = f"./output/{current_time}/"
    cam1_path = os.path.join(output_path, "cam1")
    cam2_path = os.path.join(output_path, "cam2")
    if not os.path.exists(cam1_path):
        os.makedirs(cam1_path)
    if not os.path.exists(cam2_path):
        os.makedirs(cam2_path)

    # 初始化红外相机
    infrared_camera = InfraredCamera()
    if infrared_camera.initialize():
        infrared_camera.start()
        try:
            frame_count = 0  # 用于生成唯一的文件名
            while True:
                ir_data1 = infrared_camera.get_ir_data(device_index=0)
                ir_data2 = infrared_camera.get_ir_data(device_index=1)
                if ir_data1 is None or ir_data2 is None:
                    continue

                # 转换为 RGB 图像
                ir_image1 = cv2.cvtColor(ir_data1, cv2.COLOR_GRAY2RGB)
                ir_image2 = cv2.cvtColor(ir_data2, cv2.COLOR_GRAY2RGB)

                # 显示图像
                cv2.imshow("Infrared Device1", ir_image1)
                cv2.imshow("Infrared Device2", ir_image2)

                # 检测按键
                key = cv2.waitKey(1)
                if key != -1:  # 如果有按键按下
                    # 保存图像
                    frame_count += 1
                    filename1 = os.path.join(cam1_path, f"frame_{frame_count}.png")
                    filename2 = os.path.join(cam2_path, f"frame_{frame_count}.png")
                    cv2.imwrite(filename1, ir_image1)
                    cv2.imwrite(filename2, ir_image2)
                    print(f"Saved frame {frame_count} to {filename1} and {filename2}")

        except KeyboardInterrupt:
            pass
        finally:
            infrared_camera.stop()